package com.timeriver.feature_project

import org.apache.spark.ml.feature.PolynomialExpansion
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
  * 多项式扩展是将你的特征扩展到多项式空间的过程，该空间是由原始维度的n度组合制定的。
  *   两个特征项：x、y
  *   最高次幂为2：x^2、y^2、xy、x、y
  *   最高次幂为3：x^3、y^3、x^2y、xy^2、x^2、y^2、xy、x、y
  */
object PolynomialExpansionDemo {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .master("local[6]")
      .appName("多项式扩展")
      .getOrCreate()

    val vectors = Array(
      Vectors.dense(12, 2),
      Vectors.dense(15, 6),
      Vectors.dense(18, 4)
    )

    val data: DataFrame = session.createDataFrame(vectors.map(Tuple1.apply)).toDF("features")

    val expansion: PolynomialExpansion = new PolynomialExpansion()
      .setInputCol("features")
      .setOutputCol("polyFeatures")
      /** 多项式最高次幂 */
      .setDegree(2)

    val res: DataFrame = expansion.transform(data)
    res.show(false)
  }
}
